package com.xz.courseschedule.ai;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;

/**
 * @Author: zys
 * @CreateTime: 2025-05-29
 */


@Service
public class OllamaService {

    private final WebClient webClient;

    public OllamaService(WebClient.Builder webClientBuilder) {
        this.webClient = webClientBuilder.baseUrl("http://localhost:11434").build();
    }

    public Flux<String> queryModel(String prompt) {  // 修改返回类型为Flux
        ObjectMapper mapper = new ObjectMapper();
        try {
            String requestBody = mapper.writeValueAsString(Map.of(
                    "model", "deepseek-r1:7b",
                    "prompt", prompt,
                    "stream", true  // 保持启用流式
            ));
            
            return webClient.post()
                    .uri("/api/generate")
                    .body(Mono.just(requestBody), String.class)
                    .retrieve()
                    .bodyToFlux(String.class);  // 返回Flux流而不是阻塞获取结果
        } catch (JsonProcessingException e) {
            return Flux.error(e);  // 异常时返回错误流
        }
    }
}